/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 性能评估器相关源文件
 * Author: sunzhepeng
 * Create: 2022-8-1
 */

#include "hockney_model.h"
#include "evaluator_log_def.h"

namespace HcclEvaluator {

HockneyModel::HockneyModel(HockneyModelParam param) : alpha_(param.alpha), beta_(param.beta), gamma_(param.gamma)
{
}

float HockneyModel::CalcCost(u32 nodeNum, u64 byteNum, SubgraphOpType op, AlgType alg)
{
    if (nodeNum <= 1 || byteNum == 0) {
        return 0;
    }

    // 一对多，多对一等输入输出count不一致的算子，耗时计算中统一使用总count，算子输入中的为单rank上的count，此处做转换
    if (op == SubgraphOpType::REDUCE_SCATTER || op == SubgraphOpType::ALLGATHER ||
        op == SubgraphOpType::SCATTER || op == SubgraphOpType::GATHER) {
        byteNum *= nodeNum;
    }

    float cost = -1;
    switch (alg) {
        case AlgType::PAIR_WISE:
            cost = CalcPairWiseCost(nodeNum, byteNum, op);
            break;
        case AlgType::RING:
            cost = CalcRingCost(nodeNum, byteNum, op);
            break;
        case AlgType::MULTI_RING:
            cost = CalcRingCost(nodeNum, byteNum / 4, op);  // 910环数为4，均分数据量，后期考虑传入环数
            break;
        case AlgType::MESH:
            cost = CalcMeshCost(nodeNum, byteNum, op);
            break;
        case AlgType::BINARY_HD:
            EVALUATOR_INFO("No Support for BINARY_HD, use RECURSIVE_HD instead\n");
        case AlgType::RECURSIVE_HD:
            cost = CalcHalvingDoublingCost(nodeNum, byteNum, op);
            break;
        default:
            EVALUATOR_WARNING("No Support algType\n");
    }
    return cost;
}

bool HockneyModel::CheckPower2(u32 num)
{
    return (num & (num - 1)) == 0;
}

float HockneyModel::CalcLog2(u32 num)
{
    return log(num) / log(2);  // 计算以2为底的对数
}

float HockneyModel::CalcFloorLog2(u32 num)
{
    return floor(CalcLog2(num));
}

float HockneyModel::CalcCeilLog2(u32 num)
{
    return ceil(CalcLog2(num));
}

float HockneyModel::CalcFloorPower2(u32 num)
{
    return pow(2, CalcFloorLog2(num));  // 计算向下最接近的2的整数次幂
}

float HockneyModel::CalcPairWiseCost(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    float cost = -1;

    switch (op) {
        case SubgraphOpType::REDUCE_SCATTER:
            EVALUATOR_WARNING("reducescatter not support pairwise algorithm\n");
            break;
        case SubgraphOpType::ALLGATHER:
            EVALUATOR_WARNING("allgather not support pairwise algorithm\n");
            break;
        case SubgraphOpType::BROADCAST:
            EVALUATOR_WARNING("broadcast not support pairwise algorithm\n");
            break;
        case SubgraphOpType::ALLREDUCE:
            EVALUATOR_WARNING("allreduce not support pairwise algorithm\n");
            break;
        case SubgraphOpType::GATHER:
            EVALUATOR_WARNING("gather not support pairwise algorithm\n");
            break;
        case SubgraphOpType::SCATTER:
            EVALUATOR_WARNING("scatter not support pairwise algorithm\n");
            break;
        case SubgraphOpType::REDUCE:
            EVALUATOR_WARNING("reduce not support pairwise algorithm\n");
            break;
        case SubgraphOpType::ALLTOALL:
            if (nodeNum == 0) {
                HCCL_ERROR("[Calc][PairWiseCost]In Calc PairWiseCost the value of nodeNum is 0");
                return cost;
            }
            cost = (nodeNum - 1) * alpha_ + (nodeNum - 1.0f) / nodeNum * byteNum * beta_;
            break;
        default:
            EVALUATOR_WARNING("No Support match of algType and opType\n");
    }
    return cost;
}

float HockneyModel::CalcRingCost(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    float nBeta = byteNum * beta_;
    float nGamma = byteNum * gamma_;
    if (nodeNum == 0) {
        HCCL_ERROR("[Calc][RingCost]In Calc RingCost the value of nodeNum is 0");
        return 0.0;
    }
    float coeff = (nodeNum - 1.0f) / nodeNum;
    float cost = -1;

    switch (op) {
        case SubgraphOpType::REDUCE_SCATTER:
            cost = (nodeNum - 1) * alpha_ + coeff * nBeta + coeff * nGamma;
            break;
        case SubgraphOpType::ALLGATHER:
            cost = (nodeNum - 1) * alpha_ + coeff * nBeta;
            break;
        case SubgraphOpType::BROADCAST:
            // Broadcast = Scatter + Allgather
            cost = 2 * (nodeNum - 1) * alpha_ + 2 * coeff * nBeta;  // 2(p-1)α + 2(p-1)nβ/p
            break;
        case SubgraphOpType::ALLREDUCE:
            // Allreduce = ReduceScatter + Allgather
            cost = 2 * (nodeNum - 1) * alpha_ + 2 * coeff * nBeta + coeff * nGamma; // 2(p-1)α + 2(p-1)nβ/p + (p-1)nγ/p
            break;
        case SubgraphOpType::GATHER:
            cost = (nodeNum - 1) * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::SCATTER:
            cost = (nodeNum - 1) * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::REDUCE:
            // Reduce = ReduceScatter + Gather
            cost = 2 * (nodeNum - 1) * alpha_ + 2 * coeff * nBeta + coeff * nGamma; // 2(p-1)α + 2(p-1)nβ/p + (p-1)nγ/p
            break;
        case SubgraphOpType::ALLTOALL:
            EVALUATOR_WARNING("all2all not support ring algorithm\n");
            break;
        default:
            EVALUATOR_WARNING("No Support match of algType and opType\n");
    }
    return cost;
}

float HockneyModel::CalcMeshCost(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    float nBeta = byteNum * beta_;
    float nGamma = byteNum * gamma_;
    float cost = -1;

    switch (op) {
        case SubgraphOpType::REDUCE_SCATTER:
            if (nodeNum == 0) {
                HCCL_ERROR("[Calc][MeshCost]In Calc MeshCost the value of nodeNum is 0");
                return cost;
            }
            cost = (nodeNum - 1) * alpha_ + 1.0f / nodeNum * nBeta + 1.0f / nodeNum * nGamma;
            break;
        case SubgraphOpType::ALLGATHER:
            cost = (nodeNum - 1) * alpha_ + 1.0f / nodeNum * nBeta;
            break;
        case SubgraphOpType::BROADCAST:
            // Broadcast = Scatter + Allgather
            cost = 2 * (nodeNum - 1) * alpha_ + 2.0f / nodeNum * nBeta;  // 2(p-1)α + (2.0f)nβ/p
            break;
        case SubgraphOpType::ALLREDUCE:
            // Allreduce = ReduceScatter + Allgather
            // 2(p-1)α+(2.0f)nβ/p+(1.0)nγ/p
            cost = 2 * (nodeNum - 1) * alpha_ + 2.0f / nodeNum * nBeta + 1.0f / nodeNum * nGamma;
            break;
        case SubgraphOpType::GATHER:
            cost = (nodeNum - 1) * alpha_ + 1.0f / nodeNum * nBeta ;
            break;
        case SubgraphOpType::SCATTER:
            cost = (nodeNum - 1) * alpha_ + 1.0f / nodeNum * nBeta ;
            break;
        case SubgraphOpType::REDUCE:
            // Reduce = ReduceScatter + Gather
            // 2(p-1)α+(2.0f)nβ/p+(1.0)nγ/p
            cost = 2 * (nodeNum - 1) * alpha_ + 2.0f / nodeNum * nBeta + 1.0f / nodeNum * nGamma;
            break;
        case SubgraphOpType::ALLTOALL:
            EVALUATOR_WARNING("all2all not support mesh algorithm\n");
            break;
        default:
            EVALUATOR_WARNING("No Support match of algType and opType\n");
    }
    return cost;
}

float HockneyModel::CalcHalvingDoublingCost(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    if (CheckPower2(nodeNum)) {
        return CalcHalvingDoublingCostPower2(nodeNum, byteNum, op);
    } else {
        return CalcHalvingDoublingCostNonPower2(nodeNum, byteNum, op);
    }
}

float HockneyModel::CalcHalvingDoublingCostPower2(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    float nBeta = byteNum * beta_;
    float nGamma = byteNum * gamma_;
    if (nodeNum == 0) {
        HCCL_ERROR("[Calc][HalvingDoublingCostPower2]In Calc HalvingDoublingCostPower2 the value of nodeNum is 0");
        return 0.0;
    }
    float coeff = (nodeNum - 1.0f) / nodeNum;
    float cost = -1;

    float log2Value = CalcLog2(nodeNum); // 计算nodeNum以2为底的对数
    switch (op) {
        case SubgraphOpType::REDUCE_SCATTER:
            cost = log2Value * alpha_ + coeff * nBeta + coeff * nGamma;
            break;
        case SubgraphOpType::ALLGATHER:
            cost = log2Value * alpha_ + coeff * nBeta;
            break;
        case SubgraphOpType::BROADCAST:
            cost = log2Value * alpha_ + log2Value * nBeta;
            break;
        case SubgraphOpType::ALLREDUCE:
            // Allreduce = ReduceScatter + Allgather
            cost = 2 * log2Value * alpha_ + 2 * coeff * nBeta + coeff * nGamma; // 2(logp)α + 2nβ(p-1)/p + nγ(p-1)/p
            break;
        case SubgraphOpType::GATHER:
            cost = log2Value * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::SCATTER:
            cost = log2Value * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::REDUCE:
            // Reduce = ReduceScatter + Gather
            cost = 2 * log2Value * alpha_ + 2 * coeff * nBeta + coeff * nGamma; // 2(logp)α + 2nβ(p-1)/p + nγ(p-1)/p
            break;
        case SubgraphOpType::ALLTOALL:
            EVALUATOR_WARNING("all2all not support halving-doubling algorithm\n");
            break;
        default:
            EVALUATOR_WARNING("No Support match of algType and opType\n");
    }
    return cost;
}

float HockneyModel::CalcHalvingDoublingCostNonPower2(u32 nodeNum, u64 byteNum, SubgraphOpType op)
{
    float nBeta = byteNum * beta_;
    float nGamma = byteNum * gamma_;
    if (nodeNum == 0) {
        HCCL_ERROR("[Calc][HalvingDoublingCostNonPower2]In Calc HalvingDoublingCostNonPower2 " \
            "the value of nodeNum is 0");
        return 0.0;
    }
    float coeff = (nodeNum - 1.0f) / nodeNum;
    float cost = -1;

    auto nobalanceHD = [=](u32 nodeNum) {
        float k = CalcFloorLog2(nodeNum);
        float sum = 0;
        for (float i = 1; i <= k; i++) {
            sum += ceil(nodeNum / pow(2, i));
        }
        return sum / nodeNum;
    };

    float tmp;
    float ceilLog2;
    float floorPower2;
    switch (op) {
        case SubgraphOpType::REDUCE_SCATTER:
            tmp = nobalanceHD(nodeNum) + 1;
            // 参见SD文档，2为recursive方法在非2的整数次幂的额外发送，1.0f / p 为额外数据量
            cost = (CalcFloorLog2(nodeNum) + 2) * alpha_ + (tmp + 1.0f / nodeNum) * nBeta + tmp * nGamma;
            break;
        case SubgraphOpType::ALLGATHER:
            tmp = nobalanceHD(nodeNum) + 1;
            // 参见SD文档，2为recursive方法在非2的整数次幂的额外发送，1.0f / p 为额外数据量
            cost = (CalcFloorLog2(nodeNum) + 2) * alpha_ + (tmp + 1.0f / nodeNum) * nBeta;
            break;
        case SubgraphOpType::BROADCAST:
            ceilLog2 = CalcCeilLog2(nodeNum);
            cost = ceilLog2 * alpha_ + ceilLog2 * nBeta;
            break;
        case SubgraphOpType::ALLREDUCE:
            // Allreduce = ReduceScatter + Allgather
            floorPower2 = CalcFloorPower2(nodeNum);
            tmp = (floorPower2 - 1) / floorPower2 + 1;
            // 参见SD文档，2为recursive方法在非2的整数次幂的额外发送
            cost = 2 * (CalcFloorLog2(nodeNum) + 1) * alpha_ + 2 * tmp * nBeta + tmp * nGamma;
            break;
        case SubgraphOpType::GATHER:
            ceilLog2 = CalcCeilLog2(nodeNum);
            cost = ceilLog2 * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::SCATTER:
            ceilLog2 = CalcCeilLog2(nodeNum);
            cost = ceilLog2 * alpha_ + coeff * nBeta ;
            break;
        case SubgraphOpType::REDUCE:
            // Reduce = ReduceScatter + Gather
            floorPower2 = CalcFloorPower2(nodeNum);
            tmp = (floorPower2 - 1) / floorPower2 ;
            // 参见SD文档，2为recursive方法在非2的整数次幂的额外发送
            cost = (2 * CalcFloorLog2(nodeNum) + 1) * alpha_ + (2 * tmp + 1) * nBeta + (tmp + 1) * nGamma;
            break;
        case SubgraphOpType::ALLTOALL:
            EVALUATOR_WARNING("all2all not support halving-doubling algorithm\n");
            break;
        default:
            EVALUATOR_WARNING("No Support match of algType and opType\n");
    }
    return cost;
}

}